#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>

#include "MSTK.h"


#ifdef __cplusplus
extern "C" {
#endif

  /* 
     This function partition the mesh into submeshes.

     GlobalID, GEntID, and GEntDim are retained in submesh for
     vertex and element.
     
     MasterParID is assigned for both global and submeshes.
5B     
     Caller should allocate memory for submeshes: 
     Mesh_ptr *submeshes = (Mesh_ptr*) MSTK_malloc(num*sizeof(Mesh_ptr);
  */
int        MESH_Surf_Partition(Mesh_ptr mesh, Mesh_ptr *submeshes, int num);
int        MESH_Vol_Partition(Mesh_ptr mesh, Mesh_ptr *submeshes, int num);

int MESH_Partition(Mesh_ptr mesh, Mesh_ptr *submeshes, int num) {
  int i, j, nf, nr, idx, natt, ncomp;
  MVertex_ptr global_mv;
  MFace_ptr global_mf;
  MRegion_ptr global_mr;
  char attname[256];
  MType mtype;
  MAttType att_type;
  MAttrib_ptr attrib;

  /* basic mesh information */
  nf = MESH_Num_Faces(mesh);
  nr = MESH_Num_Regions(mesh);

  /* assign global ID equal to ID for global mesh*/
  idx = 0;
  while((global_mv = MESH_Next_Vertex(mesh,&idx)))
    MV_Set_GlobalID(global_mv,MV_ID(global_mv));

  if (nr) {
    idx = 0;  
    while((global_mr = MESH_Next_Region(mesh,&idx)))
      MR_Set_GlobalID(global_mr,MR_ID(global_mr));
    MESH_Vol_Partition(mesh,submeshes,num);
  }
  else if(nf) {
    idx = 0;
    while((global_mf = MESH_Next_Face(mesh,&idx)))
      MF_Set_GlobalID(global_mf,MF_ID(global_mf));
    MESH_Surf_Partition(mesh,submeshes,num);
  }
  else {
    MSTK_Report("MESH_Partition()","This is not a valid mstk file for partition",ERROR);
    exit(-1);
  }


  /* assign attrib list */
  natt = MESH_Num_Attribs(mesh);
  for(j = 0; j < natt; j++) {
    attrib = MESH_Attrib(mesh,j);
    MAttrib_Get_Name(attrib,attname);
    att_type = MAttrib_Get_Type(attrib);
    ncomp = MAttrib_Get_NumComps(attrib);
    mtype = MAttrib_Get_EntDim(attrib);
    for(i = 0; i < num; i++) {
      if(ncomp == 1)
	attrib =  MAttrib_New(submeshes[i], attname, att_type, mtype);
      else
	attrib =  MAttrib_New(submeshes[i], attname, att_type, mtype, ncomp);
    }
  }

  return 1;
}

int MESH_Surf_Partition(Mesh_ptr mesh, Mesh_ptr *submeshes, int num) {
  int i, j, idx, max_nfv;
  int nv, nf, nfv;
  MVertex_ptr mv, vertex, global_mv;
  MFace_ptr mf, global_mf;
  List_ptr mfverts;

  /* basic mesh information */
  RepType rtype = MESH_RepType(mesh);
  nv = MESH_Num_Vertices(mesh);
  nf = MESH_Num_Faces(mesh);

  /* Partition the mesh */
  int *part = (int *)MSTK_malloc(nf*sizeof(int));
  int ok = MESH_PartitionWithMetis(mesh, num, &part);
  
  int part_no;
  /* global to local vertex ID mapping */
  int *v_g2l = (int *)MSTK_malloc(nv*sizeof(int));
  int index_vertex, index_face;
  /* List of faces and vertices of this partition */
  List_ptr part_vertices, part_faces;
  /* each face belongs only to one partition, so no need to Mark  */
  int mkvid = MSTK_GetMarker();
  /* for setting the partition id */
  int global_mkvid = MSTK_GetMarker();

  /* for each partition */
  for(part_no = 0; part_no < num; part_no++) {
    max_nfv = 0;
    part_vertices = List_New(10), part_faces = List_New(10);
    index_vertex = 0, index_face = 0;
    idx = 0, i = 0;
    while((mf = MESH_Next_Face(mesh,&idx))) {
      if(part[i++] != part_no)
	continue;
      /* add faces */
      index_face++;
      List_Add(part_faces,mf);
      
      /* add vertices */
      mfverts = MF_Vertices(mf,1,0);
      nfv = List_Num_Entries(mfverts);
      if(nfv > max_nfv)
	max_nfv = nfv;
      for(j = 0; j < nfv; j++) {
	mv = List_Entry(mfverts,j);
	if (!MEnt_IsMarked(mv,mkvid)) {
	  v_g2l[MEnt_ID(mv)-1] = index_vertex++;
	  List_Add(part_vertices,mv);
	  MEnt_Mark(mv,mkvid);
	}
      }
      List_Delete(mfverts);
    } 

    
    /* collect the coordinates and store the global vertex ID */
    int *v_l2g = (int *)MSTK_malloc(index_vertex*sizeof(int)); 
    typedef double (*COORDS)[3];
    COORDS xyz =(COORDS) MSTK_malloc(index_vertex*3*sizeof(double));
    for (i = 0; i < index_vertex; i++) {
      vertex = List_Entry(part_vertices,i);
      MV_Coords(vertex,xyz[i]);
      v_l2g[i] = MEnt_ID(vertex);
    }
    /* collect the face-vertex connectivity and store the global face ID*/ 
    int *f_l2g = (int *)MSTK_malloc(index_face*sizeof(int)); 
    int *nfvs = (int *)MSTK_malloc(index_face*sizeof(int));
    int *fvids = (int *)MSTK_malloc(index_face*max_nfv*sizeof(int));
    for (i = 0; i < index_face; i++) {
      mf = List_Entry(part_faces,i);
      mfverts = MF_Vertices(mf,1,0);
      nfv = List_Num_Entries(mfverts);
      nfvs[i] = nfv;
      f_l2g[i] = MEnt_ID(mf);
      for (j = 0; j< nfv; j++) {
	mv = List_Entry(mfverts,j);
	fvids[i*max_nfv+j] = v_g2l[MEnt_ID(mv)-1];
      }
      List_Delete(mfverts);
    }
    
    
    /* build the submesh */
    Mesh_ptr submesh = MESH_New(rtype);
    MESH_InitFromGenDesc(submesh,index_vertex,xyz,index_face,nfvs,&fvids,0,NULL,NULL,NULL,NULL);

    /* recover the vertex global ID and store the local ID into master local id*/
    for(i = 0; i < index_vertex; i++) {
      mv = MESH_Vertex(submesh,i);
      /*for mpi*/
      global_mv = MESH_VertexFromID(mesh,v_l2g[i]);
      /* copy GEnt info */
      MV_Set_GEntID(mv,MV_GEntID(global_mv));
      MV_Set_GEntDim(mv,MV_GEntDim(global_mv));
      /*if master par id is not assigned*/
      if (!MEnt_IsMarked(global_mv,global_mkvid)) {
	MEnt_Mark(global_mv,global_mkvid);
	/* set master partition id in the global mesh */
	MV_Set_MasterParID(global_mv,part_no);
	/* local is the same as global */
	MV_Set_MasterParID(mv,part_no);
	MV_Set_GlobalID(mv,v_l2g[i]);
      }
      /* if global is assigned, set local equal to global*/
      else {
	MV_Set_MasterParID(mv,MV_MasterParID(global_mv));
	MV_Set_GlobalID(mv,MV_GlobalID(global_mv));
      }	
      /*end for mpi*/
    }


    /* recover the face global ID and store the local ID into master local id */
    for(i = 0; i < index_face; i++) {
      mf = MESH_Face(submesh,i);
      /* for mpi */
      global_mf = MESH_FaceFromID(mesh,f_l2g[i]);
      /* copy GEnt info */
      MF_Set_GEntID(mf,MF_GEntID(global_mf));
      MF_Set_GEntDim(mf,MF_GEntDim(global_mf));
      /*set global face */
      MF_Set_MasterParID(global_mf,part_no);
      /*set local face*/
      MF_Set_MasterParID(mf,part_no);
      MF_Set_GlobalID(mf,f_l2g[i]);
      /*end for mpi*/
    }

    submeshes[part_no] = submesh;     
    
    /* clear the marker */
    List_Unmark(part_vertices,mkvid); 
    List_Delete(part_vertices);
    List_Delete(part_faces);
   
    MSTK_free(xyz);
    MSTK_free(nfvs);
    MSTK_free(fvids);
    MSTK_free(v_l2g);
    MSTK_free(f_l2g);
  }

  List_Unmark(MESH_Vertex_List(mesh),global_mkvid);
  MSTK_FreeMarker(global_mkvid);
  MSTK_FreeMarker(mkvid);
  MSTK_free(v_g2l);
  MSTK_free(part);
  
  return ok;
}


int MESH_Vol_Partition(Mesh_ptr mesh, Mesh_ptr *submeshes, int num) {
  int i, j, idx, max_nrv;
  int nv, nr, nrv;
  MVertex_ptr mv, vertex, global_mv;
  MRegion_ptr mr, global_mr;
  List_ptr mrverts;

  /* basic mesh information */
  RepType rtype = MESH_RepType(mesh);
  nv = MESH_Num_Vertices(mesh);
  nr = MESH_Num_Regions(mesh);

  /* Partition the mesh */
  int *part = (int *)MSTK_malloc(nr*sizeof(int));
  int ok = MESH_PartitionWithMetis(mesh, num, &part);
  
  int part_no;
  /* global to local vertex ID mapping */
  int *v_g2l = (int *)MSTK_malloc(nv*sizeof(int));
  int index_vertex, index_region;
  /* List of regions and vertices of this partition */
  List_ptr part_vertices, part_regions;
  /* each region belongs only to one partition, so no need to Mark  */
  int mkvid = MSTK_GetMarker();
  /* for setting the partition id */
  int global_mkvid = MSTK_GetMarker();

  /* for each partition */
  for(part_no = 0; part_no < num; part_no++) {
    max_nrv = 0;
    part_vertices = List_New(10), part_regions = List_New(10);
    index_vertex = 0, index_region = 0;
    idx = 0, i = 0;
    while((mr = MESH_Next_Region(mesh,&idx))) {
      if(part[i++] != part_no)
	continue;
      /* add regions */
      index_region++;
      List_Add(part_regions,mr);
      
      /* add vertices */
      mrverts = MR_Vertices(mr);
      nrv = List_Num_Entries(mrverts);
      if(nrv > max_nrv)
	max_nrv = nrv;
      for(j = 0; j < nrv; j++) {
	mv = List_Entry(mrverts,j);
	if (!MEnt_IsMarked(mv,mkvid)) {
	  v_g2l[MEnt_ID(mv)-1] = index_vertex++;
	  List_Add(part_vertices,mv);
	  MEnt_Mark(mv,mkvid);
	}
      }
      List_Delete(mrverts);
    } 

    
    /* collect the coordinates and store the global vertex ID */
    int *v_l2g = (int *)MSTK_malloc(index_vertex*sizeof(int)); 
    typedef double (*COORDS)[3];
    COORDS xyz =(COORDS) MSTK_malloc(index_vertex*3*sizeof(double));
    for (i = 0; i < index_vertex; i++) {
      vertex = List_Entry(part_vertices,i);
      MV_Coords(vertex,xyz[i]);
      v_l2g[i] = MEnt_ID(vertex);
    }
    /* collect the region-vertex connectivity and store the global region ID*/ 
    int *r_l2g = (int *)MSTK_malloc(index_region*sizeof(int)); 
    int *nrvs = (int *)MSTK_malloc(index_region*sizeof(int));
    int *rvids = (int *)MSTK_malloc(index_region*max_nrv*sizeof(int));
    for (i = 0; i < index_region; i++) {
      mr = List_Entry(part_regions,i);
      mrverts = MR_Vertices(mr);
      nrv = List_Num_Entries(mrverts);
      nrvs[i] = nrv;
      r_l2g[i] = MEnt_ID(mr);
      for (j = 0; j< nrv; j++) {
	mv = List_Entry(mrverts,j);
	rvids[i*max_nrv+j] = v_g2l[MEnt_ID(mv)-1];
      }
      List_Delete(mrverts);
    }
    
    
    /* build the submesh */
    Mesh_ptr submesh = MESH_New(rtype);
    MESH_InitFromGenDesc(submesh,index_vertex,xyz,0,NULL,NULL,index_region,nrvs,&rvids,NULL,NULL);

    /* recover the vertex global ID and store the local ID into master local id*/
    for(i = 0; i < index_vertex; i++) {
      mv = MESH_Vertex(submesh,i);
      /*for mpi*/
      global_mv = MESH_VertexFromID(mesh,v_l2g[i]);
      /* copy GEnt info */
      MV_Set_GEntID(mv,MV_GEntID(global_mv));
      MV_Set_GEntDim(mv,MV_GEntDim(global_mv));
      /*if master par id is not assigned*/
      if (!MEnt_IsMarked(global_mv,global_mkvid)) {
	MEnt_Mark(global_mv,global_mkvid);
	/* set master partition id in the global mesh */
	MV_Set_MasterParID(global_mv,part_no);
	/* local is the same as global */
	MV_Set_MasterParID(mv,part_no);
	MV_Set_GlobalID(mv,v_l2g[i]);
      }
      /* if global is assigned, set local equal to global*/
      else {
	MV_Set_MasterParID(mv,MV_MasterParID(global_mv));
	MV_Set_GlobalID(mv,MV_GlobalID(global_mv));
      }	
      /*end for mpi*/
    }


    /* recover the region global ID and store the local ID into master local id */
    for(i = 0; i < index_region; i++) {
      mr = MESH_Region(submesh,i);
      /* for mpi */
      global_mr = MESH_RegionFromID(mesh,r_l2g[i]);
      /* copy GEnt info */
      MR_Set_GEntID(mr,MR_GEntID(global_mr));
      MR_Set_GEntDim(mr,MR_GEntDim(global_mr));
      /*set global region */
      MR_Set_MasterParID(global_mr,part_no);
      /*set local region*/
      MR_Set_MasterParID(mr,part_no);
      MR_Set_GlobalID(mr,r_l2g[i]);
      /*end for mpi*/
    }

    submeshes[part_no] = submesh;     
    
    /* clear the marker */
    List_Unmark(part_vertices,mkvid); 
    List_Delete(part_vertices);
    List_Delete(part_regions);
   
    MSTK_free(xyz);
    MSTK_free(nrvs);
    MSTK_free(rvids);
    MSTK_free(v_l2g);
    MSTK_free(r_l2g);
  }

  List_Unmark(MESH_Vertex_List(mesh),global_mkvid);
  MSTK_FreeMarker(global_mkvid);
  MSTK_FreeMarker(mkvid);
  MSTK_free(v_g2l);
  MSTK_free(part);
  
  return ok;
}

  
#ifdef __cplusplus
}
#endif

